import matplotlib.pyplot as plt
import matplotlib.animation as animation
import numpy as np
import time


# 插入排序算法
def insertion_sort(arr):
    n = len(arr)
    steps = []

    for i in range(1, n):
        key = arr[i]
        j = i - 1
        while j >= 0 and key < arr[j]:
            steps.append((arr.copy(), j, j + 1))
            arr[j + 1] = arr[j]
            steps.append((arr.copy(), j, j + 1))  # 记录排序过程和当前比较的位置
            j -= 1
        steps.append((arr.copy(), j + 1, j + 1))
        arr[j + 1] = key
        steps.append((arr.copy(), j + 1, j + 1))  # 记录排序过程和当前比较的位置

    return steps


# 随机生成一个数组
arr = np.random.randint(0, 100, 20)

# 获取排序步骤
steps = insertion_sort(arr.copy())

# 创建动画
fig, ax = plt.subplots()
ax.set_xlim(-1, len(arr))
ax.set_ylim(0, max(arr) + 10)

bar_rects = ax.bar(range(len(arr)), arr, align="edge")


def animate(i):
    current_step = steps[i][0]
    compare_idx1 = steps[i][1]
    compare_idx2 = steps[i][2]

    for idx, (rect, val) in enumerate(zip(bar_rects, current_step)):
        if idx == compare_idx1 or idx == compare_idx2:
            time.sleep(0.05)
            rect.set_color('r')  # 将比较的柱子设为红色
        else:
            rect.set_color('b')  # 其余柱子设为蓝色
        rect.set_height(val)

    return bar_rects


ani = animation.FuncAnimation(fig, animate, frames=len(steps), interval=50, blit=True)

plt.show()
